//! 亲和性和反亲和性约束
//! 
//! 提供灵活的沙箱放置约束机制

use std::collections::HashMap;
use serde::{Serialize, Deserialize};

/// 亲和性类型
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AffinityType {
    /// 节点亲和性
    NodeAffinity,
    /// 沙箱亲和性
    SandboxAffinity,
    /// 区域亲和性
    ZoneAffinity,
}

/// 亲和性约束集合
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AffinityConstraints {
    /// 节点亲和性规则
    pub node_affinity: Option<NodeAffinity>,
    /// 沙箱亲和性规则
    pub sandbox_affinity: Option<SandboxAffinity>,
    /// 节点反亲和性规则
    pub node_anti_affinity: Option<NodeAntiAffinity>,
    /// 沙箱反亲和性规则
    pub sandbox_anti_affinity: Option<SandboxAntiAffinity>,
}

/// 节点亲和性
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeAffinity {
    /// 必需的选择器（硬约束）
    pub required_during_scheduling_ignored_during_execution: Vec<NodeSelectorTerm>,
    /// 偏好的选择器（软约束）
    pub preferred_during_scheduling_ignored_during_execution: Vec<PreferredSchedulingTerm>,
}

/// 沙箱亲和性
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SandboxAffinity {
    /// 必需的选择器
    pub required_during_scheduling_ignored_during_execution: Vec<SandboxAffinityTerm>,
    /// 偏好的选择器
    pub preferred_during_scheduling_ignored_during_execution: Vec<WeightedSandboxAffinityTerm>,
}

/// 节点反亲和性
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeAntiAffinity {
    /// 必需的选择器（硬约束）
    pub required_during_scheduling_ignored_during_execution: Vec<NodeSelectorTerm>,
    /// 偏好的选择器（软约束）
    pub preferred_during_scheduling_ignored_during_execution: Vec<PreferredSchedulingTerm>,
}

impl NodeAntiAffinity {
    /// 创建新的节点反亲和性配置
    pub fn new() -> Self {
        Self {
            required_during_scheduling_ignored_during_execution: Vec::new(),
            preferred_during_scheduling_ignored_during_execution: Vec::new(),
        }
    }
}

/// 沙箱反亲和性
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SandboxAntiAffinity {
    /// 必需的选择器
    pub required_during_scheduling_ignored_during_execution: Vec<SandboxAffinityTerm>,
    /// 偏好的选择器
    pub preferred_during_scheduling_ignored_during_execution: Vec<WeightedSandboxAffinityTerm>,
}

/// 节点选择器术语
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeSelectorTerm {
    /// 匹配表达式
    pub match_expressions: Vec<NodeSelectorRequirement>,
    /// 匹配字段
    pub match_fields: Vec<NodeSelectorRequirement>,
}

/// 节点选择器要求
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeSelectorRequirement {
    /// 键
    pub key: String,
    /// 操作符
    pub operator: SelectorOperator,
    /// 值列表
    pub values: Vec<String>,
}

/// 选择器操作符
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum SelectorOperator {
    /// 在列表中
    In,
    /// 不在列表中
    NotIn,
    /// 存在键
    Exists,
    /// 不存在键
    DoesNotExist,
    /// 大于
    Gt,
    /// 小于
    Lt,
}

/// 偏好调度术语
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PreferredSchedulingTerm {
    /// 权重（1-100）
    pub weight: i32,
    /// 偏好条件
    pub preference: NodeSelectorTerm,
}

/// 沙箱亲和性术语
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SandboxAffinityTerm {
    /// 标签选择器
    pub label_selector: LabelSelector,
    /// 命名空间
    pub namespaces: Vec<String>,
    /// 拓扑键
    pub topology_key: String,
}

/// 加权沙箱亲和性术语
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WeightedSandboxAffinityTerm {
    /// 权重（1-100）
    pub weight: i32,
    /// 沙箱亲和性术语
    pub sandbox_affinity_term: SandboxAffinityTerm,
}

/// 标签选择器
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LabelSelector {
    /// 匹配标签
    pub match_labels: HashMap<String, String>,
    /// 匹配表达式
    pub match_expressions: Vec<LabelSelectorRequirement>,
}

/// 标签选择器要求
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LabelSelectorRequirement {
    /// 键
    pub key: String,
    /// 操作符
    pub operator: SelectorOperator,
    /// 值列表
    pub values: Vec<String>,
}

/// 亲和性规则
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AffinityRule {
    /// 规则名称
    pub name: String,
    /// 规则类型
    pub affinity_type: AffinityType,
    /// 是否为硬约束
    pub hard: bool,
    /// 权重（软约束时使用）
    pub weight: Option<f64>,
    /// 选择器
    pub selector: AffinitySelector,
}

/// 亲和性选择器
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AffinitySelector {
    /// 节点选择器
    Node(NodeSelectorTerm),
    /// 沙箱选择器
    Sandbox(SandboxAffinityTerm),
    /// 标签选择器
    Label(LabelSelector),
}

impl AffinityConstraints {
    /// 创建新的亲和性约束
    pub fn new() -> Self {
        Self {
            node_affinity: None,
            sandbox_affinity: None,
            node_anti_affinity: None,
            sandbox_anti_affinity: None,
        }
    }

    /// 添加节点亲和性
    pub fn with_node_affinity(mut self, affinity: NodeAffinity) -> Self {
        self.node_affinity = Some(affinity);
        self
    }

    /// 添加沙箱亲和性
    pub fn with_sandbox_affinity(mut self, affinity: SandboxAffinity) -> Self {
        self.sandbox_affinity = Some(affinity);
        self
    }

    /// 添加节点反亲和性
    pub fn with_node_anti_affinity(mut self, anti_affinity: NodeAntiAffinity) -> Self {
        self.node_anti_affinity = Some(anti_affinity);
        self
    }

    /// 添加沙箱反亲和性
    pub fn with_sandbox_anti_affinity(mut self, anti_affinity: SandboxAntiAffinity) -> Self {
        self.sandbox_anti_affinity = Some(anti_affinity);
        self
    }

    /// 检查约束是否为空
    pub fn is_empty(&self) -> bool {
        self.node_affinity.is_none()
            && self.sandbox_affinity.is_none()
            && self.node_anti_affinity.is_none()
            && self.sandbox_anti_affinity.is_none()
    }

    /// 验证约束的合法性
    pub fn validate(&self) -> Result<(), String> {
        // 验证权重范围
        if let Some(node_affinity) = &self.node_affinity {
            for term in &node_affinity.preferred_during_scheduling_ignored_during_execution {
                if term.weight < 1 || term.weight > 100 {
                    return Err(format!("节点亲和性权重必须在1-100之间，当前: {}", term.weight));
                }
            }
        }

        if let Some(sandbox_affinity) = &self.sandbox_affinity {
            for term in &sandbox_affinity.preferred_during_scheduling_ignored_during_execution {
                if term.weight < 1 || term.weight > 100 {
                    return Err(format!("沙箱亲和性权重必须在1-100之间，当前: {}", term.weight));
                }
            }
        }

        Ok(())
    }
}

impl Default for AffinityConstraints {
    fn default() -> Self {
        Self::new()
    }
}

impl NodeAffinity {
    /// 创建新的节点亲和性
    pub fn new() -> Self {
        Self {
            required_during_scheduling_ignored_during_execution: Vec::new(),
            preferred_during_scheduling_ignored_during_execution: Vec::new(),
        }
    }

    /// 添加必需的选择器
    pub fn add_required_selector(mut self, term: NodeSelectorTerm) -> Self {
        self.required_during_scheduling_ignored_during_execution.push(term);
        self
    }

    /// 添加偏好的选择器
    pub fn add_preferred_selector(mut self, term: PreferredSchedulingTerm) -> Self {
        self.preferred_during_scheduling_ignored_during_execution.push(term);
        self
    }
}

impl SandboxAffinity {
    /// 创建新的沙箱亲和性
    pub fn new() -> Self {
        Self {
            required_during_scheduling_ignored_during_execution: Vec::new(),
            preferred_during_scheduling_ignored_during_execution: Vec::new(),
        }
    }

    /// 添加必需的选择器
    pub fn add_required_selector(mut self, term: SandboxAffinityTerm) -> Self {
        self.required_during_scheduling_ignored_during_execution.push(term);
        self
    }

    /// 添加偏好的选择器
    pub fn add_preferred_selector(mut self, term: WeightedSandboxAffinityTerm) -> Self {
        self.preferred_during_scheduling_ignored_during_execution.push(term);
        self
    }
}

impl NodeSelectorTerm {
    /// 创建新的节点选择器术语
    pub fn new() -> Self {
        Self {
            match_expressions: Vec::new(),
            match_fields: Vec::new(),
        }
    }

    /// 添加匹配表达式
    pub fn add_match_expression(mut self, requirement: NodeSelectorRequirement) -> Self {
        self.match_expressions.push(requirement);
        self
    }

    /// 添加匹配字段
    pub fn add_match_field(mut self, requirement: NodeSelectorRequirement) -> Self {
        self.match_fields.push(requirement);
        self
    }
}

impl NodeSelectorRequirement {
    /// 创建新的节点选择器要求
    pub fn new(key: String, operator: SelectorOperator, values: Vec<String>) -> Self {
        Self { key, operator, values }
    }

    /// 检查是否匹配给定的标签
    pub fn matches(&self, labels: &HashMap<String, String>) -> bool {
        let label_value = labels.get(&self.key);

        match self.operator {
            SelectorOperator::In => {
                label_value.map_or(false, |v| self.values.contains(v))
            }
            SelectorOperator::NotIn => {
                label_value.map_or(true, |v| !self.values.contains(v))
            }
            SelectorOperator::Exists => {
                label_value.is_some()
            }
            SelectorOperator::DoesNotExist => {
                label_value.is_none()
            }
            SelectorOperator::Gt => {
                if let (Some(label_val), Some(compare_val)) = (label_value, self.values.first()) {
                    match (label_val.parse::<f64>(), compare_val.parse::<f64>()) {
                        (Ok(lv), Ok(cv)) => lv > cv,
                        _ => false,
                    }
                } else {
                    false
                }
            }
            SelectorOperator::Lt => {
                if let (Some(label_val), Some(compare_val)) = (label_value, self.values.first()) {
                    match (label_val.parse::<f64>(), compare_val.parse::<f64>()) {
                        (Ok(lv), Ok(cv)) => lv < cv,
                        _ => false,
                    }
                } else {
                    false
                }
            }
        }
    }
}

impl LabelSelector {
    /// 创建新的标签选择器
    pub fn new() -> Self {
        Self {
            match_labels: HashMap::new(),
            match_expressions: Vec::new(),
        }
    }

    /// 添加匹配标签
    pub fn add_match_label(mut self, key: String, value: String) -> Self {
        self.match_labels.insert(key, value);
        self
    }

    /// 添加匹配表达式
    pub fn add_match_expression(mut self, requirement: LabelSelectorRequirement) -> Self {
        self.match_expressions.push(requirement);
        self
    }

    /// 检查是否匹配给定的标签
    pub fn matches(&self, labels: &HashMap<String, String>) -> bool {
        // 检查匹配标签
        for (key, value) in &self.match_labels {
            if labels.get(key) != Some(value) {
                return false;
            }
        }

        // 检查匹配表达式
        for expr in &self.match_expressions {
            let requirement = NodeSelectorRequirement {
                key: expr.key.clone(),
                operator: expr.operator.clone(),
                values: expr.values.clone(),
            };
            if !requirement.matches(labels) {
                return false;
            }
        }

        true
    }
}

impl Default for LabelSelector {
    fn default() -> Self {
        Self::new()
    }
}

/// 亲和性约束构建器
pub struct AffinityConstraintsBuilder {
    constraints: AffinityConstraints,
}

impl AffinityConstraintsBuilder {
    /// 创建新的构建器
    pub fn new() -> Self {
        Self {
            constraints: AffinityConstraints::new(),
        }
    }

    /// 添加区域亲和性（常用的便利方法）
    pub fn prefer_zone(mut self, zone: String, weight: i32) -> Self {
        let node_affinity = self.constraints.node_affinity.get_or_insert_with(NodeAffinity::new);
        
        let requirement = NodeSelectorRequirement::new(
            "zone".to_string(),
            SelectorOperator::In,
            vec![zone],
        );
        
        let term = NodeSelectorTerm::new().add_match_expression(requirement);
        let preferred_term = PreferredSchedulingTerm { weight, preference: term };
        
        node_affinity.preferred_during_scheduling_ignored_during_execution.push(preferred_term);
        self
    }

    /// 添加实例类型偏好
    pub fn prefer_instance_type(mut self, instance_type: String, weight: i32) -> Self {
        let node_affinity = self.constraints.node_affinity.get_or_insert_with(NodeAffinity::new);
        
        let requirement = NodeSelectorRequirement::new(
            "instance-type".to_string(),
            SelectorOperator::In,
            vec![instance_type],
        );
        
        let term = NodeSelectorTerm::new().add_match_expression(requirement);
        let preferred_term = PreferredSchedulingTerm { weight, preference: term };
        
        node_affinity.preferred_during_scheduling_ignored_during_execution.push(preferred_term);
        self
    }

    /// 避免某个区域（反亲和性）
    pub fn avoid_zone(mut self, zone: String) -> Self {
        let node_anti_affinity = self.constraints.node_anti_affinity.get_or_insert_with(NodeAntiAffinity::new);
        
        let requirement = NodeSelectorRequirement::new(
            "zone".to_string(),
            SelectorOperator::In,
            vec![zone],
        );
        
        let term = NodeSelectorTerm::new().add_match_expression(requirement);
        node_anti_affinity.required_during_scheduling_ignored_during_execution.push(term);
        self
    }

    /// 构建约束
    pub fn build(self) -> Result<AffinityConstraints, String> {
        self.constraints.validate()?;
        Ok(self.constraints)
    }
}

impl Default for AffinityConstraintsBuilder {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_node_selector_requirement_in() {
        let requirement = NodeSelectorRequirement::new(
            "zone".to_string(),
            SelectorOperator::In,
            vec!["us-west-1a".to_string(), "us-west-1b".to_string()],
        );

        let mut labels = HashMap::new();
        labels.insert("zone".to_string(), "us-west-1a".to_string());
        
        assert!(requirement.matches(&labels));

        labels.insert("zone".to_string(), "us-east-1a".to_string());
        assert!(!requirement.matches(&labels));
    }

    #[test]
    fn test_node_selector_requirement_not_in() {
        let requirement = NodeSelectorRequirement::new(
            "zone".to_string(),
            SelectorOperator::NotIn,
            vec!["us-west-1a".to_string()],
        );

        let mut labels = HashMap::new();
        labels.insert("zone".to_string(), "us-west-1b".to_string());
        
        assert!(requirement.matches(&labels));

        labels.insert("zone".to_string(), "us-west-1a".to_string());
        assert!(!requirement.matches(&labels));
    }

    #[test]
    fn test_node_selector_requirement_exists() {
        let requirement = NodeSelectorRequirement::new(
            "gpu".to_string(),
            SelectorOperator::Exists,
            vec![],
        );

        let mut labels = HashMap::new();
        labels.insert("gpu".to_string(), "nvidia-k80".to_string());
        
        assert!(requirement.matches(&labels));

        labels.remove("gpu");
        assert!(!requirement.matches(&labels));
    }

    #[test]
    fn test_node_selector_requirement_gt() {
        let requirement = NodeSelectorRequirement::new(
            "memory".to_string(),
            SelectorOperator::Gt,
            vec!["8192".to_string()],
        );

        let mut labels = HashMap::new();
        labels.insert("memory".to_string(), "16384".to_string());
        
        assert!(requirement.matches(&labels));

        labels.insert("memory".to_string(), "4096".to_string());
        assert!(!requirement.matches(&labels));
    }

    #[test]
    fn test_label_selector() {
        let selector = LabelSelector::new()
            .add_match_label("zone".to_string(), "us-west-1a".to_string())
            .add_match_expression(LabelSelectorRequirement {
                key: "instance-type".to_string(),
                operator: SelectorOperator::In,
                values: vec!["m5.large".to_string(), "m5.xlarge".to_string()],
            });

        let mut labels = HashMap::new();
        labels.insert("zone".to_string(), "us-west-1a".to_string());
        labels.insert("instance-type".to_string(), "m5.large".to_string());
        
        assert!(selector.matches(&labels));

        labels.insert("zone".to_string(), "us-west-1b".to_string());
        assert!(!selector.matches(&labels));
    }

    #[test]
    fn test_affinity_constraints_builder() {
        let constraints = AffinityConstraintsBuilder::new()
            .prefer_zone("us-west-1a".to_string(), 80)
            .prefer_instance_type("m5.large".to_string(), 60)
            .avoid_zone("us-east-1a".to_string())
            .build()
            .unwrap();

        assert!(constraints.node_affinity.is_some());
        assert!(constraints.node_anti_affinity.is_some());
        
        let node_affinity = constraints.node_affinity.as_ref().unwrap();
        assert_eq!(node_affinity.preferred_during_scheduling_ignored_during_execution.len(), 2);
        
        let anti_affinity = constraints.node_anti_affinity.as_ref().unwrap();
        assert_eq!(anti_affinity.required_during_scheduling_ignored_during_execution.len(), 1);
    }

    #[test]
    fn test_affinity_constraints_validation() {
        let mut constraints = AffinityConstraints::new();
        constraints.node_affinity = Some(NodeAffinity {
            required_during_scheduling_ignored_during_execution: Vec::new(),
            preferred_during_scheduling_ignored_during_execution: vec![
                PreferredSchedulingTerm {
                    weight: 150, // 无效权重
                    preference: NodeSelectorTerm::new(),
                }
            ],
        });

        assert!(constraints.validate().is_err());

        // 修正权重
        if let Some(ref mut node_affinity) = constraints.node_affinity {
            node_affinity.preferred_during_scheduling_ignored_during_execution[0].weight = 80;
        }

        assert!(constraints.validate().is_ok());
    }

    #[test]
    fn test_empty_constraints() {
        let constraints = AffinityConstraints::new();
        assert!(constraints.is_empty());

        let constraints = AffinityConstraints::new()
            .with_node_affinity(NodeAffinity::new());
        assert!(!constraints.is_empty());
    }
}